#!/usr/bin/env python
import sqlite3
import os
import time
from math import *
from turtle import color
from matplotlib import pyplot as plt

Timestep = 0.1

# Control output of results
def listControl(list,num):
    endPoint = len(list)-num
    for i in range(endPoint, len(list)):
        list[i] = list[endPoint]

# Plot figures after getting result database
def display(IsSave, T_3,T_RP,T_RV,T_PV,  E_RP,E_PV,E_RV,  U_RP,U_PV,U_RV,  TD):
    if IsSave:
        time_now = time.strftime("%Y%m%d-%H:%M", time.localtime())
        pic_save_path = '/home/ach/nzx_实验/shujv'+time_now
        os.makedirs(pic_save_path)

    db_file = '/home/ach/nzx_workspace/src/webots_ros/scripts/DT_data.db'
    if os.path.isfile(db_file):    
        conn = sqlite3.connect(db_file)
        cursor = conn.cursor()
        show_sql0 = 'select PE_x,PE_y,VE_x,VE_y,Ref_x,Ref_y,PE_th,PE_vx,PE_vy,VE_th,VE_vx,VE_vy,VE_vth,Ref_th,PE_vth, PE_ux, PE_uy, PE_uth, VE_ux, VE_uy, VE_uth, DT_PE_x, DT_PE_y from DT_states'
        dis_res = cursor.execute(show_sql0)
        all_data = dis_res.fetchall()
        all_data = all_data[2:-20]
       
        PE_x_list = []
        PE_y_list = []
        PE_th_list = []
        VE_x_list = []
        VE_y_list = []
        VE_th_list = []
        Ref_x_list = []
        Ref_y_list = []
        Ref_th_list = []
        PE_vx_list = []
        PE_vy_list = []
        PE_vth_list = []
        VE_vx_list = []
        VE_vy_list = []
        VE_vth_list = []
        PE_ux_list = []
        PE_uy_list = []
        PE_uth_list = []
        VE_ux_list = []
        VE_uy_list = []
        VE_uth_list = []

        PE_xDT_list = []
        PE_yDT_list = []
        
        # error between Ref and PE
        PE_eth_list = []
        PE_ex_list = []
        PE_ey_list = []

        # error between VE and PE
        VE_eth_list = []
        VE_ex_list = []
        VE_ey_list = []

        # error between Ref and VE
        RV_eth_list = []
        RV_ex_list = []
        RV_ey_list = []

        # error between PE_pose and DT
        PE_DT_ex = []
        PE_DT_ey = []

        for row in all_data:
            PE_x_list.append(row[0])
            PE_y_list.append(row[1])
            VE_x_list.append(row[2])
            VE_y_list.append(row[3])
            Ref_x_list.append(row[4])
            Ref_y_list.append(row[5])
            PE_th_list.append(row[6])
            PE_vx_list.append(row[7])
            PE_vy_list.append(row[8])
            VE_th_list.append(row[9])
            VE_vx_list.append(row[10])
            VE_vy_list.append(row[11])
            VE_vth_list.append(row[12])
            Ref_th_list.append(row[13])
            PE_vth_list.append(row[14])
            PE_ux_list.append(row[15])
            PE_uy_list.append(row[16])
            PE_uth_list.append(row[17])
            VE_ux_list.append(row[18])
            VE_uy_list.append(row[19])
            VE_uth_list.append(row[20])

            PE_xDT_list.append(row[21])
            PE_yDT_list.append(row[22])
        
        # listControl(PE_x_list,34)
        # listControl(PE_y_list,34)
        # listControl(VE_x_list,30)
        # listControl(VE_y_list,30)

        T_list = []
        for i in range(0,len(Ref_x_list)):
            T_list.append(Timestep*i)

        for i in range(0,len(Ref_x_list)-0):
            
            ee_th = Ref_th_list[i]-PE_th_list[i]
            EE_th = PE_th_list[i]-VE_th_list[i]
            RV_e_th = Ref_th_list[i]-VE_th_list[i]

            if ee_th >= pi:
                ee_th -= 2*pi
            if ee_th <= -pi:
                ee_th += 2*pi
            if EE_th >= 2*pi:
                EE_th -= 2*pi
            if EE_th <= -2*pi:
                EE_th += 2*pi
            if RV_e_th >= pi:
                RV_e_th -= 2*pi
            if RV_e_th <= -pi:
                RV_e_th += 2*pi

            if ee_th >= 0.15:
                ee_th = 0

            PE_eth_list.append(ee_th)
            PE_ex_list.append(Ref_x_list[i]-PE_x_list[i])
            PE_ey_list.append(Ref_y_list[i]-PE_y_list[i])

            VE_eth_list.append(EE_th)
            VE_ex_list.append(PE_x_list[i]-VE_x_list[i])
            VE_ey_list.append(PE_y_list[i]-VE_y_list[i])

            RV_eth_list.append(RV_e_th)
            RV_ex_list.append(Ref_x_list[i]-VE_x_list[i])
            RV_ey_list.append(Ref_y_list[i]-VE_y_list[i])

            PE_DT_ex.append(PE_x_list[i]-PE_xDT_list[i])
            PE_DT_ey.append(PE_y_list[i]-PE_yDT_list[i])
        
        # Trajectories
        if T_3: 
            plt.figure(num='Track of Ref, PE and VE', dpi=300)
            manager = plt.get_current_fig_manager()
            manager.resize(*manager.window.maxsize())
            plt.plot(Ref_x_list, Ref_y_list, '--', color='gold', lw='2',zorder=1,label="Reference")
            plt.scatter(Ref_x_list[0], Ref_y_list[0],s=100,zorder=2,color='gold',marker='*')
            plt.plot(PE_x_list, PE_y_list,zorder=3,label="PE",color='darkorange', lw='2')
            plt.scatter(PE_x_list[0], PE_y_list[0],s=100,zorder=4,color='darkorange',marker='*')
            plt.plot(VE_x_list, VE_y_list,zorder=5,label="VE",color='steelblue', lw='2')
            plt.scatter(VE_x_list[0], VE_y_list[0],s=100,zorder=6,color='steelblue',marker='*')
            plt.scatter(Ref_x_list[-1], Ref_y_list[-1],s=150,zorder=7,color='gold',marker='.')
            plt.scatter(PE_x_list[-1], PE_y_list[-1],s=150,zorder=8,color='darkorange',marker='.')
            plt.scatter(VE_x_list[-1], VE_y_list[-1],s=150,zorder=9,color='steelblue',marker='.')
            # plt.xlim([0,4.8])
            # plt.ylim([0,3.6])
            plt.xlabel(r'$X (m)$')
            plt.ylabel(r'$Y (m)$')
            plt.legend(loc='best')
            
            if IsSave:
                plt.savefig(pic_save_path+"/Track of Ref, PE and VE.svg", dpi=300)
            
        # Ref & PE trajectories
        if T_RP:
            plt.figure(num='Track of Ref and PE', dpi=250)
            manager = plt.get_current_fig_manager()
            manager.resize(*manager.window.maxsize())
            plt.plot(Ref_x_list, Ref_y_list, color='gold', lw='2',zorder=1,label="Reference")
            plt.scatter(Ref_x_list[0], Ref_y_list[0],s=100,zorder=2,color='gold',marker='*')
            plt.plot(PE_x_list, PE_y_list,zorder=3,label="PE",color='darkorange', lw='2')
            plt.scatter(PE_x_list[0], PE_y_list[0],s=100,zorder=4,color='darkorange',marker='*')
            plt.scatter(Ref_x_list[-1], Ref_y_list[-1],s=150,zorder=7,color='gold',marker='.')
            plt.scatter(PE_x_list[-1], PE_y_list[-1],s=150,zorder=8,color='darkorange',marker='.')
            # plt.xlim([0,4.8])
            # plt.ylim([0,3.6])
            plt.xlabel(r'$X (m)$')
            plt.ylabel(r'$Y (m)$')
            plt.legend(loc='best')
            
            if IsSave:
                plt.savefig(pic_save_path+"/Track of Ref and PE.svg", dpi=300)
            
        # PE & VE trajectories
        if T_PV:
            plt.figure(num='Track of PE and VE', dpi=300)
            manager = plt.get_current_fig_manager()
            manager.resize(*manager.window.maxsize())
            plt.plot(PE_x_list, PE_y_list,zorder=3,label="PE",color='darkorange', lw='2')
            plt.scatter(PE_x_list[0], PE_y_list[0],s=100,zorder=4,color='darkorange',marker='*')
            plt.plot(VE_x_list, VE_y_list,zorder=5,label="VE",color='steelblue', lw='2')
            plt.scatter(VE_x_list[0], VE_y_list[0],s=100,zorder=6,color='steelblue',marker='*')
            plt.scatter(PE_x_list[-1], PE_y_list[-1],s=150,zorder=8,color='darkorange',marker='.')
            plt.scatter(VE_x_list[-1], VE_y_list[-1],s=150,zorder=9,color='steelblue',marker='.')
            # plt.xlim([0,4.8])
            # plt.ylim([0,3.6])
            plt.xlabel(r'$X (m)$')
            plt.ylabel(r'$Y (m)$')
            plt.legend(loc='best')

            if IsSave:
                plt.savefig(pic_save_path+"/Track of PE and VE.svg", dpi=300)

        # Ref & VE trajectories
        if T_RV:
            plt.figure(num='Track of Ref and VE', dpi=300)
            manager = plt.get_current_fig_manager()
            manager.resize(*manager.window.maxsize())
            # plt.plot(Ref_x_list, Ref_y_list,  color='gold', lw='2',label="Reference")
            # plt.scatter(Ref_x_list[0], Ref_y_list[0],s=100,zorder=2,color='gold',marker='*')
            # plt.plot(VE_x_list, VE_y_list,label="VE",color='steelblue', lw='2')
            plt.plot(Ref_x_list, Ref_y_list,  color='gold', lw='2',zorder=1,label="Reference")
            plt.scatter(Ref_x_list[0], Ref_y_list[0],s=100,zorder=2,color='gold',marker='*')
            plt.plot(VE_x_list, VE_y_list,zorder=5,label="VE",color='steelblue', lw='2')
            plt.scatter(VE_x_list[0], VE_y_list[0],s=100,zorder=6,color='steelblue',marker='*')
            plt.scatter(Ref_x_list[-1], Ref_y_list[-1],s=150,zorder=7,color='gold',marker='.')
            plt.scatter(VE_x_list[-1], VE_y_list[-1],s=150,zorder=9,color='steelblue',marker='.')
            # plt.xlim([0,4.8])
            # plt.ylim([0,3.6])
            plt.xlabel(r'$X (m)$')
            plt.ylabel(r'$Y (m)$')
            plt.legend(loc='best')

            if IsSave:
                plt.savefig(pic_save_path+"/Track of Ref and VE.svg", dpi=300)

        # Ref & PE error
        if E_RP:
            plt.figure('Error between Ref and PE', dpi=300)
            manager = plt.get_current_fig_manager()
            manager.resize(*manager.window.maxsize())
            plt.subplots_adjust(hspace=0.6)
            plt.subplot(311)
            plt.plot(T_list,PE_ex_list,lw = '2')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$e_{xPE} (m)$')
            plt.subplot(312)
            plt.plot(T_list,PE_ey_list,lw = '2')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$e_{yPE} (m)$')
            plt.subplot(313)
            plt.plot(T_list,PE_eth_list,lw = '2')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$e_{\Theta PE} (rad)$')

            if IsSave:
                plt.savefig(pic_save_path+"/Error between Ref and PE.svg", dpi=300)

        # PE & VE error
        if E_PV:
            plt.figure('Error between PE and VE', dpi=300)
            manager = plt.get_current_fig_manager()
            manager.resize(*manager.window.maxsize())
            plt.subplots_adjust(hspace=0.6)
            plt.subplot(311)
            plt.plot(T_list,VE_ex_list,lw = '2')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$e_{xVE} (m)$')
            plt.subplot(312)
            plt.plot(T_list,VE_ey_list,lw = '2')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$e_{yVE} (m)$')
            plt.subplot(313)
            plt.plot(T_list,VE_eth_list,lw = '2')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$e_{\Theta VE} (rad)$')

            if IsSave:
                plt.savefig(pic_save_path+"/Error between PE and VE.svg", dpi=300)

        # Ref & VE error
        if E_RV:
            plt.figure('Error between Ref and VE', dpi=300)
            manager = plt.get_current_fig_manager()
            manager.resize(*manager.window.maxsize())
            plt.subplots_adjust(hspace=0.6)
            plt.subplot(311)
            plt.plot(T_list,RV_ex_list,lw = '1')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$E_{xVE} (m)$')
            plt.subplot(312)
            plt.plot(T_list,RV_ey_list,lw = '1')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$E_{yVE} (m)$')
            plt.subplot(313)
            plt.plot(T_list,RV_eth_list,lw = '1')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$E_{\Theta VE} (rad)$')

            if IsSave:
                plt.savefig(pic_save_path+"/Error between Ref and VE.svg", dpi=300)

        # Ref & PE control law
        if U_RP:
            plt.figure('U_PE', dpi=300)
            manager = plt.get_current_fig_manager()
            manager.resize(*manager.window.maxsize())
            plt.subplots_adjust(hspace=0.6)
            plt.subplot(311)
            plt.plot(T_list,PE_ux_list,lw = '2')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$u_{xPE} (m/s)$')
            plt.subplot(312)
            plt.plot(T_list,PE_uy_list,lw = '2')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$u_{yPE} (m/s)$')
            plt.subplot(313)
            plt.plot(T_list,PE_uth_list,lw = '2')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$u_{\theta PE} (rad/s)$')

            if IsSave:
                plt.savefig(pic_save_path+"/U_PE.svg", dpi=300)
        
        # PE & VE control law
        if U_PV:
            plt.figure('U_VE', dpi=300)
            manager = plt.get_current_fig_manager()
            manager.resize(*manager.window.maxsize())
            plt.subplots_adjust(hspace=0.6)
            plt.subplot(311)
            plt.plot(T_list,VE_ux_list,lw = '2')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$u_{xVE} (m/s)$')
            plt.subplot(312)
            plt.plot(T_list,VE_uy_list,lw = '2')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$u_{yVE} (m/s)$')
            plt.subplot(313)
            plt.plot(T_list,VE_uth_list,lw = '2')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$u_{\theta VE} (rad/s)$')
            if IsSave:
                plt.savefig(pic_save_path+"/U_VE.svg", dpi=300)

        # Ref & VE control law
        if U_RV:
            plt.figure('U_RV', dpi=300)
            manager = plt.get_current_fig_manager()
            manager.resize(*manager.window.maxsize())
            plt.subplots_adjust(hspace=0.6)
            plt.subplot(311)
            plt.plot(T_list,VE_ux_list,lw = '1')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$u_{xPE} (m/s)$')
            plt.subplot(312)
            plt.plot(T_list,VE_uy_list,lw = '1')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$u_{yPE} (m/s)$')
            plt.subplot(313)
            plt.plot(T_list,VE_uth_list,lw = '1')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$u_{\theta PE} (rad/s)$')
        
        # TD
        if TD:
            plt.figure('eDT', dpi=300)
            manager = plt.get_current_fig_manager()
            manager.resize(*manager.window.maxsize())
            plt.subplots_adjust(hspace=0.6)
            plt.subplot(211)
            plt.plot(T_list,PE_DT_ex,lw = '2')
            plt.subplot(212)
            plt.plot(T_list,PE_DT_ey,lw = '2')
            # plt.scatter(PE_xDT_list[0], 0,s=100)

            plt.figure('DT', dpi=300)
            manager = plt.get_current_fig_manager()
            manager.window.showMaximized()
            plt.subplots_adjust(hspace=0.6)
            plt.subplot(211)
            plt.plot(T_list,PE_x_list, ls='dashed',lw = '2')
            plt.plot(T_list,PE_xDT_list,lw = '1')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$x_{PE} (m)$')
            plt.subplot(212)
            plt.plot(T_list,PE_y_list, ls='dashed',lw = '2')
            plt.plot(T_list,PE_yDT_list,lw = '1')
            plt.xlabel('Time (s)')
            plt.ylabel(r'$y_{PE} (m)$')  

        # # PE 与 Ref angle
        # plt.figure('angle of Ref and PE', dpi=200)
        # plt.subplots_adjust(hspace=0.6)
        # plt.subplot(211)
        # plt.plot(T_list,Ref_th_list,lw = '2')
        # plt.subplot(212)
        # plt.plot(T_list,PE_th_list,lw = '2')
        # plt.show()

        # PE velocity
        # plt.figure('PE Velocity', dpi=200)
        # plt.subplots_adjust(hspace=0.6)
        # plt.subplot(311)
        # plt.plot(T_list,PE_vx_list,lw = '2')
        # plt.xlabel('Time (s)')
        # plt.ylabel(r'$v_{xPE} (m)$')
        # plt.subplot(312)
        # plt.plot(T_list,PE_vy_list,lw = '2')
        # plt.xlabel('Time (s)')
        # plt.ylabel(r'$v_{yPE} (m)$')
        # plt.subplot(313)
        # plt.plot(T_list,PE_vth_list,lw = '2')
        # plt.xlabel('Time (s)')
        # plt.ylabel(r'$v_{\omega PE} (rad)$')
        # plt.show()

        # PE control laws
        # plt.figure('U_PE', dpi=200)
        # plt.subplots_adjust(hspace=0.6)
        # plt.subplot(311)
        # plt.plot(T_list,PE_ux_list,lw = '2')
        # plt.xlabel('Time (s)')
        # plt.ylabel(r'$u_{xPE} (m)$')
        # plt.subplot(312)
        # plt.plot(T_list,PE_uy_list,lw = '2')
        # plt.xlabel('Time (s)')
        # plt.ylabel(r'$u_{yPE} (m)$')
        # plt.subplot(313)
        # plt.plot(T_list,PE_uth_list,lw = '2')
        # plt.xlabel('Time (s)')
        # plt.ylabel(r'$u_{\theta PE} (rad)$')
        # plt.show()

        # # VE velocity
        # plt.figure('VE Velocity', dpi=200)
        # plt.subplots_adjust(hspace=0.6)
        # plt.subplot(311)
        # plt.plot(T_list,VE_vx_list,lw = '2')
        # plt.xlabel('Time (s)')
        # plt.ylabel(r'$\mu_{VE} (m)$')
        # plt.subplot(312)
        # plt.plot(T_list,VE_vy_list,lw = '2')
        # plt.xlabel('Time (s)')
        # plt.ylabel(r'$\nu_{VE} (m)$')
        # plt.subplot(313)
        # plt.plot(T_list,VE_vth_list,lw = '2')
        # plt.xlabel('Time (s)')
        # plt.ylabel(r'$\varphi_{VE} (rad)$')
        # plt.show
        
        plt.show()



    